#include <opencv2/opencv.hpp>
#include <iostream>
#include <chrono>
#include "ros/ros.h"
#include "std_msgs/String.h"
#include <stdio.h>
#include <memory>
#include <sys/time.h>
#include <atomic>
#include <cv_bridge/cv_bridge.h>
#include <sensor_msgs/Image.h>

#include "opencv2/core/core.hpp"
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/imgproc/imgproc.hpp"
#include "kbt_cam/rknnPool.hpp"
#include "kbt_cam/osiris.hpp"
#include "kbt_cam/rtsp_pusher.hpp"

#include <swarm_msgs/BoundingBoxes.h>
#include <swarm_msgs/BoundingBox.h>

#include <unordered_map>
#include <vector>
#include <sensor_msgs/NavSatFix.h>

std::mutex cache_mutex; // 如果多线程使用，需加锁保护
std::unordered_map<uint64_t, std::vector<InferenceResult>> frame_cache;

// int8_t vechile_state = 0;


std::atomic<bool> running(true);
std::atomic<bool> publish_running{true};

std::mutex frame_mtx;
cv::Mat shared_frame;

// std::atomic<uint64_t> latestFrameId = {0};     // 每帧递增
// std::atomic<uint64_t> lastPushedFrameId = {0}; // 上一次压入的帧ID

constexpr size_t MAX_CACHE_SIZE = 10;  // 可根据实际内存容量调整
std::deque<std::pair<int, cv::Mat>> original_frame_cache;
std::mutex original_frame_mtx;

uint64_t no_target_num = 0;
uint64_t aim_roi_x, aim_roi_y;

// 全局或局部定义，保存预设的 6 个窗口区域
std::vector<cv::Rect> roi_windows = {
    cv::Rect(0, 0, 740, 640),     // 窗口1
    cv::Rect(590, 0, 740, 640),   // 窗口2
    cv::Rect(1180, 0, 740, 640),   // 窗口3
    cv::Rect(0, 440, 740, 640),   // 窗口4
    cv::Rect(590, 440, 740, 640), // 窗口5
    cv::Rect(1180, 440, 740, 640)  // 窗口6
};

cv::Rect make_safe_roi(int center_x, int center_y, int roi_size, int img_width, int img_height) {
    int half = roi_size / 2;

    int x = std::max(0, center_x - half);
    int y = std::max(0, center_y - half);

    // 不允许 x+width 超出图像
    if (x + roi_size > img_width)
        x = img_width - roi_size;

    if (y + roi_size > img_height)
        y = img_height - roi_size;

    // 再次确保非负（极限边界情况）
    x = std::max(0, x);
    y = std::max(0, y);

    return cv::Rect(x, y, roi_size, roi_size);
}

// 设置线程绑定到特定核心
void bind_to_core(int core_id) {
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(core_id, &cpuset);
    pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpuset);
}

void capture_and_push_thread(cv::VideoCapture &cap, rknnPool<osiris, WindowInput, InferenceResult> &camPool) {
    bind_to_core(7);  // 可选：绑定 CPU7

    cv::Mat frame;
    uint64_t currentFrameId = 0;
    uint64_t lastFrameId = 0;

    auto last_time = std::chrono::high_resolution_clock::now();
    int frame_count = 0;

    while (running) {
        if (!cap.read(frame)) {
            std::cerr << "Failed to read frame!" << std::endl;
            continue;
        }
        // frame_count++;

        // auto now = std::chrono::high_resolution_clock::now();
        // std::chrono::duration<double> elapsed = now - last_time;

        // // 每秒统计一次帧率
        // if (elapsed.count() >= 1.0) {
        //     double fps = frame_count / elapsed.count();
        //     std::cout << "[Capture] FPS: " << fps << " frames/sec" << std::endl;

        //     // 重置计数器和时间
        //     frame_count = 0;
        //     last_time = now;
        // }
        

        currentFrameId++;
        if (no_target_num < 3)
        {
            // 存储原始图像副本（用于后续结果显示）
            {
                std::lock_guard<std::mutex> orig_lock(original_frame_mtx);
                original_frame_cache.emplace_back(currentFrameId, frame.clone());

                // 控制缓存大小，避免内存膨胀
                while (original_frame_cache.size() > MAX_CACHE_SIZE) {
                    original_frame_cache.pop_front();
                }
            }

            // 确认新的识别区域，并推入推理池
            int img_width = frame.cols;
            int img_height = frame.rows;

            cv::Rect aim_roi = make_safe_roi(aim_roi_x, aim_roi_y, 640, img_width, img_height);
           
            WindowInput win_input;
            win_input.window_id = 0;
            win_input.frame_id = currentFrameId;
            win_input.window_img = frame(aim_roi).clone();  // 拷贝 ROI 区域
            win_input.target_roi_x = aim_roi.x;
            win_input.target_roi_y = aim_roi.y;

            if (camPool.put(win_input) != 0) {
                std::cerr << "camPool.put failed for roi: " << aim_roi_x << aim_roi_y << "\n";
            }

            lastFrameId = currentFrameId;
              
        }else{
            if (currentFrameId-lastFrameId < 3)
            {
                continue;
            }
            
            // 存储原始图像副本（用于后续结果显示）
            {
                std::lock_guard<std::mutex> orig_lock(original_frame_mtx);
                original_frame_cache.emplace_back(currentFrameId, frame.clone());

                // 控制缓存大小，避免内存膨胀
                while (original_frame_cache.size() > MAX_CACHE_SIZE) {
                    original_frame_cache.pop_front();
                }
            }

            // 将图像划分为多个窗口，并推入推理池
            int window_id = 0;
            for (const auto& roi : roi_windows) {
                if ((roi.x + roi.width <= frame.cols) &&
                    (roi.y + roi.height <= frame.rows)) {
                    
                    WindowInput win_input;
                    win_input.window_id = ++window_id;
                    win_input.frame_id = currentFrameId;
                    win_input.window_img = frame(roi).clone();  // 拷贝 ROI 区域
                    win_input.target_roi_x = roi.x;
                    win_input.target_roi_y = roi.y;

                    if (camPool.put(win_input) != 0) {
                        std::cerr << "camPool.put failed for roi: " << roi << "\n";
                    }
                } else {
                    std::cerr << "Invalid ROI: " << roi << " exceeds image size\n";
                }
            }
            lastFrameId = currentFrameId;

        }
    }
}


// // 采集线程函数
// void capture_thread(cv::VideoCapture &cap) {
//     bind_to_core(7);  // 采集线程绑定到 CPU7
//     cv::Mat frame;
//     auto last_time = std::chrono::high_resolution_clock::now();
//     int frame_count = 0;
//     double fps = 0.0;
//     // cv::namedWindow("new", cv::WINDOW_AUTOSIZE);
//     while (running) {
//         if (!cap.read(frame)) {
//             std::cerr << "Failed to read frame!" << std::endl;
//             continue;
//         }
//         else{
//             // cv::imshow("new", frame);
//         //     auto now = std::chrono::high_resolution_clock::now();
//         //     frame_count++;
//         // // 每秒计算一次帧率
//         //     double seconds = std::chrono::duration<double>(now - last_time).count();
//         //     if (seconds >= 0.5) {
//         //         fps = frame_count / seconds;
//         //         frame_count = 0;
//         //         last_time = now;
//         //         std::cout << "fps:" << fps << std::endl;
//         //     }
//         }
//         // cv::Mat frame = cv::imread("/home/orangepi/kbt/py-multi-demo-3588/ball.jpg");
//         {

//             // cv::Rect roi(600, 340, 640, 640);
//             // cv::Mat cropped = frame(roi);
//             std::lock_guard<std::mutex> lock(frame_mtx);
//             shared_frame = frame.clone(); // 图像数据复制
//             latestFrameId++;
//         }
//     }
// }

template<typename ModelType>
void publish_thread_func(
    rknnPool<ModelType, WindowInput, InferenceResult> &camPool,
    ros::Publisher &result_pub,
    std::atomic<bool> &publish_running
) {
    cv::namedWindow("DetectedFrame", cv::WINDOW_AUTOSIZE);

    int frames = 0;

    struct timeval time;
    gettimeofday(&time, nullptr);
    auto startTime = time.tv_sec * 1000 + time.tv_usec / 1000;
    auto beforeTime = startTime;

    const std::array<std::string, 6> kWinNames = {
        "Result1", "Result2", "Result3",
        "Result4", "Result5", "Result6"
    };

    // for (const auto& n : kWinNames) {
    //     cv::namedWindow(n, cv::WINDOW_NORMAL);
    // }
    // cv::namedWindow(n, cv::WINDOW_NORMAL);
    const int OUT_W = 720, OUT_H = 1280, FPS = 50;
    rtsp::Pusher pusher("rtsp://192.168.1.50:8554/air", OUT_W, OUT_H, FPS, 4'000'000, 50);

    while (publish_running.load()) {
        InferenceResult result;
        bool high_speed_detect = false;
        if (camPool.get(result) == 0) {
            // std_msgs::Header header;
            // header.stamp = ros::Time::now();
            // sensor_msgs::ImagePtr msg = cv_bridge::CvImage(header, "bgr8", result).toImageMsg();
            // result_pub.publish(msg);
            std::lock_guard<std::mutex> lock(cache_mutex);

            if (result.window_id == 0)
                high_speed_detect = true;
            else
                high_speed_detect = false;

            // 放入缓存
            auto& vec = frame_cache[result.frame_id];
            vec.push_back(result);

            // 如果已满6张图
            if (vec.size() == 6 || high_speed_detect) {
                // 🔍 查找所有图中置信度最大的目标
                const InferenceResult* best_result = nullptr;
                double best_prob = -1.0;
                uint64_t Current_frame_id = 0 ;
                swarm_msgs::BoundingBox bb;
            
                for (const auto& res : vec) {
                    Current_frame_id = res.frame_id;
                    if (res.has_target) {
                        if (res.prop > best_prob) {
                            best_prob = res.prop;
                            best_result = &res;
                        }
                    }
                }
                
                if (best_result) {
                    int window_id = best_result->window_id; // ∈ [1, 6]
                    int offset_x = 0, offset_y = 0;

                    // 安全取偏移量
                    if (window_id >= 1 && window_id <= 6) {
                        const cv::Rect& roi = roi_windows[window_id - 1];
                        offset_x = roi.x;
                        offset_y = roi.y;
                    }else if (window_id == 0){
                        offset_x = best_result->target_roi_x;
                        offset_y = best_result->target_roi_y;
                    }
                    // 发布该帧的最大框
                    swarm_msgs::BoundingBoxes msg;
                    msg.header.stamp = ros::Time::now();
                    msg.header.frame_id = "camera";

                   
                    bb.probability = best_result->prop;
                    bb.xmin = best_result->x1 + offset_x;
                    bb.ymin = best_result->y1 + offset_y;
                    bb.xmax = best_result->x2 + offset_x;
                    bb.ymax = best_result->y2 + offset_y;
                    // bb.id   = best_box.id;
                    // bb.class_name = best_box.class_name;
                    // bb.point = best_box.point;

                    msg.bounding_boxes.push_back(bb);
                    result_pub.publish(msg);
                    no_target_num = 0;
                }else{
                    no_target_num++;
                }
                    

                {   
                    // printf("123\n");
                    std::lock_guard<std::mutex> lock(original_frame_mtx);
                    for (auto it = original_frame_cache.rbegin(); it != original_frame_cache.rend(); ++it) {
                        if (it->first == Current_frame_id) {
                            cv::Mat raw_img = it->second.clone();

                            if(best_result){
                            // 画矩形框
                                cv::rectangle(raw_img,
                                            cv::Point(bb.xmin, bb.ymin),
                                            cv::Point(bb.xmax, bb.ymax),
                                            cv::Scalar(0, 255, 0), 2); // 绿色框，线宽 2
                                int cx = (bb.xmin + bb.xmax) / 2;
                                int cy = (bb.ymin + bb.ymax) / 2;

                                aim_roi_x = cx;
                                aim_roi_y = cy;

                                // printf("roi_x:%d,roi_y:%d\n",aim_roi_x,aim_roi_y);

                                cv::line(raw_img, cv::Point(cx, cy), cv::Point(raw_img.cols/2, raw_img.rows/2), cv::Scalar(0, 0, 255), 1); // 水平红线
                                putText(raw_img, "uav", cv::Point(bb.xmin, bb.ymin -20), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(255, 0, 0), 2);
                                // if(vechile_state == 1)
                                //     putText(raw_img, "track", cv::Point(raw_img.cols-350, raw_img.rows -100), cv::FONT_HERSHEY_SIMPLEX, 3, cv::Scalar(255, 0, 0), 3);
                                // else if (vechile_state == 2)
                                //     putText(raw_img, "attack", cv::Point(raw_img.cols-350, raw_img.rows -100), cv::FONT_HERSHEY_SIMPLEX, 3, cv::Scalar(255, 0, 0), 3);
                            }
                            // else //by zhou :save image without detect
                            // {
                            //     std::string save_dir = "/home/orangepi/pictures_yolo_horizonal";
                            
                            //     // // 获取当前时间并生成时间戳
                            //     // auto now = std::chrono::system_clock::now();
                            //     // auto time_t = std::chrono::system_clock::to_time_t(now);
                            //     // std::tm tm = *std::localtime(&time_t);
                                
                            //     // std::ostringstream timestamp;
                            //     // timestamp << std::put_time(&tm, "%Y%m%d_%H%M%S");

                            //     double rostime = ros::Time::now().toSec();

                                
                            //     // 构造完整文件路径
                            //     std::string filename = save_dir + "/" + std::to_string(rostime) + ".jpg";
                                
                            //     // 保存图像
                            //     cv::imwrite(filename, raw_img);

                            // }

                            // 计算目标中心
                            if (high_speed_detect)
                            {
                                cv::rectangle(raw_img,
                                            cv::Point(result.target_roi_x, result.target_roi_y),
                                            cv::Point(result.target_roi_x + 640, result.target_roi_y + 640),
                                            cv::Scalar(255, 255, 0), 2); // 红色框，线宽 2
                            }
                            
                            

                            // 画十字（中心到四边）
                            cv::line(raw_img, cv::Point(raw_img.cols/2, 0), cv::Point(raw_img.cols/2, raw_img.rows - 1), cv::Scalar(0, 0, 255), 1); // 垂直红线
                            cv::line(raw_img, cv::Point(0, raw_img.rows/2), cv::Point(raw_img.cols - 1, raw_img.rows/2), cv::Scalar(0, 0, 255), 1); // 水平红线

                            // // 可选：画中心点
                            // cv::circle(raw_img, cv::Point(cx, cy), 4, cv::Scalar(255, 0, 0), -1); // 蓝色中心点

                            cv::Mat Detected_Frame, rotated_img;

                            // cv::resize(raw_img, Detected_Frame, cv::Size(), 0.55, 0.55, cv::INTER_LINEAR);
                            cv::resize(raw_img, Detected_Frame, cv::Size(), 0.3, 0.3, cv::INTER_LINEAR);

                            cv::rotate(Detected_Frame, rotated_img, cv::ROTATE_90_COUNTERCLOCKWISE);

                            // cv::Rect roi_dis(960,900,1920,1080);
                            // cv::Mat Detected_Frame = raw_img(roi_dis);

                            // 显示图像
                            cv::imshow("DetectedFrame", rotated_img);
                            pusher.push(rotated_img);
                            cv::moveWindow("DetectedFrame", 0, 0);
                            cv::waitKey(1);

                            break; // 找到了就退出
                        }
                    }
                }


            // 🔁 清理该 frame_id
                frame_cache.erase(Current_frame_id);
            }
            frames++;
            if (frames % 60 == 0) {
                gettimeofday(&time, nullptr);
                auto currentTime = time.tv_sec * 1000 + time.tv_usec / 1000;
                printf("60帧内平均帧率:\t %f fps/s\n", 60.0 / float(currentTime - beforeTime) * 1000.0);
                beforeTime = currentTime;
            }

            if (cv::waitKey(1) == 'q') // 延时1毫秒,按q键退出/Press q to exit
            {
                running = false;
                cv::destroyAllWindows();
                break;
            }
        } 
        else {
            
            std::this_thread::sleep_for(std::chrono::milliseconds(1));
        }
    }
    
    gettimeofday(&time, nullptr);
    auto endTime = time.tv_sec * 1000 + time.tv_usec / 1000;

    printf("Average:\t %f fps/s\n", float(frames) / float(endTime - startTime) * 1000.0);

}

// 显式实例化模板（避免链接错误）
template void publish_thread_func<osiris>(
    rknnPool<osiris, WindowInput, InferenceResult> &,
    ros::Publisher &,
    std::atomic<bool> &
);

// void image_push_thread_func(rknnPool<osiris, WindowInput, InferenceResult> &camPool) {
//     while (running) {
//         cv::Mat input_img;
//         uint64_t currentFrameId;

//         {
//             std::lock_guard<std::mutex> lock(frame_mtx);
//             if (shared_frame.empty()) continue;

//             currentFrameId = latestFrameId.load();

//             if (currentFrameId == lastPushedFrameId.load()) {
//                 continue;
//             }

//             // 存储原始图像
//             {
//                 std::lock_guard<std::mutex> orig_lock(original_frame_mtx);
//                 original_frame_cache.emplace_back(currentFrameId, shared_frame.clone());

//                 while (original_frame_cache.size() > MAX_CACHE_SIZE) {
//                     original_frame_cache.pop_front();
//                 }
//             }

//             int window_id = 0;
//             for (const auto& roi : roi_windows) {
//                 if ((roi.x + roi.width <= shared_frame.cols) &&
//                     (roi.y + roi.height <= shared_frame.rows)) {
//                     WindowInput win_input;
//                     window_id += 1;
//                     win_input.window_id = window_id;
//                     win_input.frame_id = currentFrameId;
//                     win_input.window_img = shared_frame(roi).clone();
//                     if (camPool.put(win_input) != 0) {
//                         std::cerr << "camPool.put failed for roi: " << roi << "\n";
//                     }
//                 } else {
//                     std::cerr << "Invalid ROI: " << roi << " exceeds image size\n";
//                 }
//             }

//             lastPushedFrameId.store(currentFrameId);
//         }
//     }
// }


int main(int argc, char** argv){
    ros::init(argc, argv, "kbt_cam_node");
    ros::NodeHandle nh;
    ros::Publisher result_pub = nh.advertise<swarm_msgs::BoundingBoxes>("/yindao/capture_drone", 10);

    // ros::Subscriber vechile_state_sub = nh.subscribe("/vehicle/position/global",10,vehiclestateCallback);

    ROS_INFO("kbt_cam_node started.");
    
    /***************************************初始化摄像头开启采集线程***********************************************/
    cv::VideoCapture capture("/dev/video0", cv::CAP_V4L2);
     if (!capture.isOpened()) {
        std::cerr << "❌ 无法打开摄像头" << std::endl;
        return -1;
    }
    capture.set(cv::CAP_PROP_FOURCC, cv::VideoWriter::fourcc('M','J','P','G'));
    capture.set(cv::CAP_PROP_FRAME_WIDTH, 1920);
    capture.set(cv::CAP_PROP_FRAME_HEIGHT, 1080);
    capture.set(cv::CAP_PROP_FPS, 120);
    std::cout << "分辨率: " << capture.get(cv::CAP_PROP_FRAME_WIDTH) << "x"
              << capture.get(cv::CAP_PROP_FRAME_HEIGHT) << "\n";


    /***************************************参数初始化***********************************************/
    char *model_name = (char *)"/home/orangepi/kbt_ws/src/kbtdj_air/kbt_cam/model/1119.rknn";
    
    int threadNum = 12;
    rknnPool<osiris, WindowInput, InferenceResult> camPool(model_name, threadNum);
    if (camPool.init() != 0)
    {
        printf("rknnPool init fail!\n");
        running = false;
        // cap_thread.join();
        return -1;
    }

     std::thread cap_push_thread([&]() {
        capture_and_push_thread(capture, camPool);
    });
    
    std::thread publish_thread([&]() {
        publish_thread_func(camPool, result_pub, publish_running);
    });

    // std::thread image_push_thread([&]() {
    //     image_push_thread_func(camPool);
    // });
    


    // while (running) {
    //     cv::Mat input_img;
    //     uint64_t currentFrameId;
    
    //     {
    //         std::lock_guard<std::mutex> lock(frame_mtx);
    //         if (shared_frame.empty()) continue;

    //         currentFrameId = latestFrameId.load();

    //         // 如果当前帧已经被处理过，就跳过
    //         if (currentFrameId == lastPushedFrameId.load()) {
    //             continue;
    //         }

    //         // 存储原始图像
    //         {
    //             std::lock_guard<std::mutex> orig_lock(original_frame_mtx);
    //             original_frame_cache.emplace_back(currentFrameId, shared_frame.clone());

    //             // 控制缓存大小，避免内存膨胀
    //             while (original_frame_cache.size() > MAX_CACHE_SIZE) {
    //                 original_frame_cache.pop_front();
    //             }
    //         }

    //         // cv::imshow("img", shared_frame.clone());
    //         int window_id = 0;
    //         for (const auto& roi : roi_windows) {
    //             // 检查 ROI 是否在图像内
    //             if ((roi.x + roi.width <= shared_frame.cols) &&
    //                 (roi.y + roi.height <= shared_frame.rows)) {
    //                 WindowInput win_input;
    //                 window_id += 1;
    //                 win_input.window_id = window_id;
    //                 win_input.frame_id = currentFrameId;
    //                 win_input.window_img = shared_frame(roi).clone();
    //                 if (camPool.put(win_input) != 0) {
    //                     std::cerr << "camPool.put failed for roi: " << roi << "\n";
    //                 }
    //             } 
    //             else {
    //                 std::cerr << "Invalid ROI: " << roi << " exceeds image size\n";
    //             }
    //         }
    //         // WindowInput win_input;
    //         // win_input.window_id = window_id;
    //         // win_input.frame_id = currentFrameId;
    //         // win_input.window_img = shared_frame.clone();

    //         // if (camPool.put(win_input) != 0)
    //         //     break;


    //         lastPushedFrameId.store(currentFrameId);
    //         // // auto read = std::chrono::high_resolution_clock::now();
    //         // // auto diff = std::chrono::duration_cast<std::chrono::milliseconds>(read - start);
    //         // std::cout << "currentFrameId: " << currentFrameId << std::endl;
    //     }

    //     ros::spin();
       
        
    //     // frames++;

    //     // if (frames % 120 == 0)
    //     // {
    //     //     gettimeofday(&time, nullptr);
    //     //     auto currentTime = time.tv_sec * 1000 + time.tv_usec / 1000;
    //     //     printf("120帧内平均帧率:\t %f fps/s\n", 120.0 / float(currentTime - beforeTime) * 1000.0);
    //     //     beforeTime = currentTime;
    //     // }
    // }

    // while (true)
    // {
    //     cv::Mat img;
    //     if (camPool.get(img) != 0)
    //         break;
    //     cv::imshow("Result", img);
    //     if (cv::waitKey(1) == 'q') // 延时1毫秒,按q键退出/Press q to exit
    //         break;
    //     frames++;
    // }

    // gettimeofday(&time, nullptr);
    // auto endTime = time.tv_sec * 1000 + time.tv_usec / 1000;

    // printf("Average:\t %f fps/s\n", float(frames) / float(endTime - startTime) * 1000.0);

    /***** ROS 主线程用于处理回调 *****/
    ros::spin();  // ✅ 会持续运行并响应 vehiclestateCallback 等回调

    // 退出时 join 所有线程
    running = false;
    publish_running = false;
    cap_push_thread.join();
    publish_thread.join();
    capture.release();
    
    return 0;
}